import torch
import torchaudio
import warnings

warnings.filterwarnings('ignore')
class ft(torch.nn.Module):
    def __init__(self,**kwargs):
        super(ft,self).__init__()
        self.kwargs = kwargs
    def forward(self,waveforms):
        fts = []
        for waveform in waveforms:
            if len(waveform.shape) == 1:
                waveform = waveform.unsqueeze(0)
            # 计算声谱图
            specgram = torch.stft(waveform, n_fft=20768, hop_length=10384, center=False, return_complex=True)

            # 获取实部和虚部
            real_part = specgram[..., 0]
            imaginary_part = specgram[..., 1]

            # 构建复数张量
            complex_tensor = torch.stack((real_part, imaginary_part), dim=-1)

            # 逆变换得到声带振动
            reconstructed_waveform = torch.istft(complex_tensor, n_fft=20768, hop_length=10384, center=False)
            reconstructed_waveform = reconstructed_waveform.view(88,-1)
            fts.append(reconstructed_waveform)
        ft = torch.stack(fts)
        return ft